
#include "PiccoloFileUtil.h"
#include "PiccoloGeometryUtil.h"

#include "Graphics/PiccoloVKQueue.h"
#include "Graphics/PiccoloVKPipeline.h"
#include "Graphics/PiccoloVKRenderPass.h"
#include "Graphics/PiccoloVKCommandBuffer.h"
#include "Graphics/PiccoloVKDescriptorSet.h"
#include "Graphics/PiccoloVKImageView.h"

#include "PiccoloEntryPoint.h"
#include "Render/PiccoloRenderTarget.h"
#include "Render/PiccoloMesh.h"
#include "Render/PiccoloTexture.h"
#include "Render/PiccoloRenderer.h"

namespace Piccolo
{
    struct GlobalUBO
    {
        glm::mat4 projMatrix { 1.f };
        glm::mat4 viewMatrix { 1.f };
    };

    struct InstanceUBO
    {
        glm::mat4 modelMatrix { 1.f };
    };

    class SandBoxApp : public PiccoloApplication
    {
    protected:
        void OnConfiguration(AppSettings *appSettings) override
        {
            appSettings->width = 1360;
            appSettings->height = 768;
            appSettings->title = "SandBox";
        }

        void OnInit() override
        {
            PiccoloRenderContext* renderContext = GetAppContext()->renderContext;
            PiccoloVKDevice* device = renderContext->GetDevice();
            PiccoloVKSwapchain* swapchain = renderContext->GetSwapchain();

            // 创建颜色附件和深度附件
            std::vector<Attachment> attachments =
            {
                {
                    .format = swapchain->GetSurfaceInfo().surfaceFormat.format,
                    .loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
                    .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
                    .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
                    .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
                    .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
                    .finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
                    .usage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT
                },
                {
                    .format = device->GetSettings().depthFormat,
                    .loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
                    .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
                    .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
                    .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
                    .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
                    .finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
                    .usage = VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT
                }
            };

            std::vector<RenderSubPass> subPasses =
            {
                {
                    .colorAttachments = { 0 },
                    .depthStencilAttachments = { 1 },
                    .sampleCount = VK_SAMPLE_COUNT_4_BIT
                }
            };
            mRenderPass = std::make_shared<PiccoloVKRenderPass>(device, attachments, subPasses);
            mRenderTarget = std::make_shared<PiccoloRenderTarget>(mRenderPass.get());
            mRenderTarget->SetColorClearValue({ 0.1f, 0.2f, 0.3f, 1.f });
            mRenderTarget->SetDepthStencilClearValue({ 1, 0 });

            mRenderer = std::make_shared<PiccoloRenderer>();

            // descriptor set，和layout(set = 0, binding = 1, std140)中的binding对应
            std::vector<VkDescriptorSetLayoutBinding> descriptorSetLayoutBindings =
            {
                {
                    .binding = 0,
                    .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                    .descriptorCount = 1,
                    .stageFlags = VK_SHADER_STAGE_VERTEX_BIT
                },
                {
                    .binding = 1,
                    .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                    .descriptorCount = 1,
                    .stageFlags = VK_SHADER_STAGE_VERTEX_BIT
                },
                {
                    .binding = 2,
                    .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                    .descriptorCount = 1,
                    .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT
                },
                {
                    .binding = 3,
                    .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                    .descriptorCount = 1,
                    .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT
                }
            };
            mDescriptorSetLayout = std::make_shared<PiccoloVKDescriptorSetLayout>(device, descriptorSetLayoutBindings);
            ShaderLayout shaderLayout =
            {
                .descriptorSetLayouts = { mDescriptorSetLayout->GetDescriptorSetLayout() }
            };
            mPipelineLayout = std::make_shared<PiccoloVKPipelineLayout>(device, PICCOLO_RES_SHADER_DIR"descriptor_set.vert",
                PICCOLO_RES_SHADER_DIR"descriptor_set.frag", shaderLayout);

            std::vector<VkVertexInputBindingDescription> vertexBindings =
            {
                {
                    .binding = 0,
                    .stride = sizeof(PiccoloVertex),
                    .inputRate = VK_VERTEX_INPUT_RATE_VERTEX
                }
            };
            std::vector<VkVertexInputAttributeDescription> vertexAttributes =
            {
                {
                    .location = 0,
                    .binding = 0,
                    .format = VK_FORMAT_R32G32B32_SFLOAT,
                    .offset = offsetof(PiccoloVertex, position)
                },
                {
                    .location = 1,
                    .binding = 0,
                    .format = VK_FORMAT_R32G32_SFLOAT,
                    .offset = offsetof(PiccoloVertex, texcoord0)
                },
                {
                    .location = 2,
                    .binding = 0,
                    .format = VK_FORMAT_R32G32B32_SFLOAT,
                    .offset = offsetof(PiccoloVertex, normal)
                }
            };

            mPipeline = std::make_shared<PiccoloVKPipeline>(device, mRenderPass.get(), mPipelineLayout.get());
            mPipeline->SetVertexInputState(vertexBindings, vertexAttributes);
            mPipeline->SetInputAssemblyState(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST)->EnableDepthTest();
            mPipeline->SetDynamicState({ VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR });
            mPipeline->SetMultisampleState(VK_SAMPLE_COUNT_4_BIT, VK_FALSE);
            mPipeline->Create();

            std::vector<VkDescriptorPoolSize> poolSizes =
            {
                {
                    .type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                    .descriptorCount = 2
                },
                {
                    .type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                    .descriptorCount = 2
                }
            };
            mDescriptoPool = std::make_shared<PiccoloVKDescriptorPool>(device, 1, poolSizes);
            // 申请descriptor set，用于往里面写UBO数据，传输给GPU处理
            mDescriptorSets = mDescriptoPool->AllocateDescriptorSet(mDescriptorSetLayout.get(), 1);
            mGlobalBuffer = std::make_shared<PiccoloVKBuffer>(device, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, sizeof(GlobalUBO), nullptr, true);
            mInstanceBuffer = std::make_shared<PiccoloVKBuffer>(device, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, sizeof(InstanceUBO), nullptr, true);
            mTexture0 = std::make_shared<PiccoloTexture>(PICCOLO_RES_TEXTURE_DIR"awesomeface.png");
            mTexture1 = std::make_shared<PiccoloTexture>(PICCOLO_RES_TEXTURE_DIR"R-C.jpeg");

            mCommandBuffers = device->GetDefaultCommandPool()->AllocateCommandBuffer(swapchain->GetImages().size());

            std::vector<PiccoloVertex> vertices;
            std::vector<uint32_t> indices;
            PiccoloGeometryUtil::CreateCube(-0.3f, 0.3f, -0.3f, 0.3f, -0.3f, 0.3f, vertices, indices);
            mCubeMesh = std::make_shared<PiccoloMesh>(vertices, indices);
        }

        void OnUpdate(float deltaTime) override
        {
            PiccoloRenderContext* renderContext = GetAppContext()->renderContext;
            PiccoloVKSwapchain* swapchain = renderContext->GetSwapchain();

            float time = std::chrono::duration<float>(std::chrono::steady_clock::now() - mStartTimePoint).count();
            mInstanceUBO.modelMatrix = glm::rotate(glm::mat4(1.f), glm::radians(17.f), glm::vec3(1.f, 0.f, 0.f));
            mInstanceUBO.modelMatrix = glm::rotate(mInstanceUBO.modelMatrix, glm::radians(time * 100.f), glm::vec3(0.f, 1.f, 0.f));

            mGlobalUBO.projMatrix = glm::perspective(glm::radians(65.f), swapchain->GetWidth() * 1.f / swapchain->GetHeight(), 0.01f, 100.f);
            // 反转一下Y轴
            mGlobalUBO.projMatrix[1][1] *= -1.f;
            mGlobalUBO.viewMatrix = glm::lookAt(glm::vec3{ 0.f, 0.f, 1.5f }, glm::vec3{ 0.f, 0.f, -1.f }, glm::vec3{ 0.f, 1.f, 0.f });
        }

        void OnRender() override
        {
            PiccoloRenderContext* renderContext = GetAppContext()->renderContext;
            PiccoloVKDevice* device = renderContext->GetDevice();
            PiccoloVKSwapchain* swapchain = renderContext->GetSwapchain();

            int32_t imageIndex;
            if (mRenderer->Begin(&imageIndex))
            {
                mRenderTarget->SetExtent({ swapchain->GetWidth(), swapchain->GetHeight() });
            }

            // 开启command buffer
            VkCommandBuffer cmdBuffer = mCommandBuffers[imageIndex];
            PiccoloVKCommandPool::BeginCommandBuffer(cmdBuffer);
            mRenderTarget->Begin(cmdBuffer);
            PiccoloVKFrameBuffer* frameBuffer = mRenderTarget->GetCurrentFrameBuffer();

            mPipeline->Bind(cmdBuffer);
            VkViewport viewport =
            {
                .x = 0,
                .y = 0,
                .width = static_cast<float>(frameBuffer->GetWidth()),
                .height = static_cast<float>(frameBuffer->GetHeight()),
                .minDepth = 0,
                .maxDepth = 1
            };
            vkCmdSetViewport(cmdBuffer, 0, 1, &viewport);
            VkRect2D scissor = { {0, 0}, {frameBuffer->GetWidth(), frameBuffer->GetHeight()} };
            vkCmdSetScissor(cmdBuffer, 0, 1, &scissor);

            // 更新UBO数据到buffer中
            mGlobalBuffer->WriteData(&mGlobalUBO);
            mInstanceBuffer->WriteData(&mInstanceUBO);
            UpdateDescriptorSet();
            // 绑定descriptor set
            vkCmdBindDescriptorSets(cmdBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipelineLayout->GetPipelineLayout(), 0, 1, mDescriptorSets.data(), 0, nullptr);

            mCubeMesh->Draw(cmdBuffer);

            mRenderTarget->End(cmdBuffer);
            PiccoloVKCommandPool::EndCommandBuffer(cmdBuffer);

            mRenderer->End(imageIndex, { cmdBuffer });
        }

        void UpdateDescriptorSet()
        {
            PiccoloRenderContext* renderContext = GetAppContext()->renderContext;
            PiccoloVKDevice* device = renderContext->GetDevice();

            VkDescriptorBufferInfo globalBufferInfo =
            {
                .buffer = mGlobalBuffer->GetBuffer(),
                .offset = 0,
                .range = sizeof(mGlobalUBO)
            };
            VkDescriptorBufferInfo instanceBufferInfo =
            {
                .buffer = mInstanceBuffer->GetBuffer(),
                .offset = 0,
                .range = sizeof(mInstanceUBO)
            };
            VkDescriptorImageInfo textureImageInfo0 =
            {
                .sampler = mTexture0->GetSampler(),
                .imageView = mTexture0->GetImageView()->GetImageView(),
                .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
            };
            VkDescriptorImageInfo textureImageInfo1 =
            {
                .sampler = mTexture1->GetSampler(),
                .imageView = mTexture1->GetImageView()->GetImageView(),
                .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
            };

            VkDescriptorSet descriptorSet = mDescriptorSets[0];

            // 描述每一个要写入的binding
            std::vector<VkWriteDescriptorSet> writeDescriptorSets =
            {
                {
                    .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                    .pNext = nullptr,
                    .dstSet = descriptorSet,
                    .dstBinding = 0,
                    .dstArrayElement = 0,
                    .descriptorCount = 1,
                    .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                    .pBufferInfo = &globalBufferInfo
                },
                {
                    .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                    .pNext = nullptr,
                    .dstSet = descriptorSet,
                    .dstBinding = 1,
                    .dstArrayElement = 0,
                    .descriptorCount = 1,
                    .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                    .pBufferInfo = &instanceBufferInfo
                },
                {
                    .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                    .pNext = nullptr,
                    .dstSet = descriptorSet,
                    .dstBinding = 2,
                    .dstArrayElement = 0,
                    .descriptorCount = 1,
                    .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                    .pImageInfo = &textureImageInfo0
                },
                {
                    .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                    .pNext = nullptr,
                    .dstSet = descriptorSet,
                    .dstBinding = 3,
                    .dstArrayElement = 0,
                    .descriptorCount = 1,
                    .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                    .pImageInfo = &textureImageInfo1
                }
            };
            vkUpdateDescriptorSets(device->GetLogicDevice(), writeDescriptorSets.size(), writeDescriptorSets.data(), 0, nullptr);
        }

        void OnDestroy() override
        {
            PiccoloRenderContext* renderContext = GetAppContext()->renderContext;
            PiccoloVKDevice* device = renderContext->GetDevice();
            vkDeviceWaitIdle(device->GetLogicDevice());
            mGlobalBuffer.reset();
            mInstanceBuffer.reset();
            mTexture0.reset();
            mTexture1.reset();
            mCubeMesh.reset();
            mCommandBuffers.clear();
            mDescriptoPool.reset();
            mDescriptorSetLayout.reset();
            mPipeline.reset();
            mPipelineLayout.reset();
            mRenderTarget.reset();
            mRenderPass.reset();
            mRenderer.reset();
        }

    private:
        std::shared_ptr<PiccoloVKRenderPass> mRenderPass;
        std::shared_ptr<PiccoloRenderTarget> mRenderTarget;
        std::shared_ptr<PiccoloRenderer> mRenderer;

        std::shared_ptr<PiccoloVKDescriptorSetLayout> mDescriptorSetLayout;
        std::shared_ptr<PiccoloVKDescriptorPool> mDescriptoPool;
        std::vector<VkDescriptorSet> mDescriptorSets;

        std::shared_ptr<PiccoloVKPipelineLayout> mPipelineLayout;
        std::shared_ptr<PiccoloVKPipeline> mPipeline;
        std::vector<VkCommandBuffer> mCommandBuffers;

        std::shared_ptr<PiccoloMesh> mCubeMesh;
        GlobalUBO mGlobalUBO;
        InstanceUBO mInstanceUBO;
        std::shared_ptr<PiccoloVKBuffer> mGlobalBuffer;
        std::shared_ptr<PiccoloVKBuffer> mInstanceBuffer;
        std::shared_ptr<PiccoloTexture> mTexture0;
        std::shared_ptr<PiccoloTexture> mTexture1;
    };
}

Piccolo::PiccoloApplication* CreateApplicationEntryPoint()
{
    return new Piccolo::SandBoxApp();
}